#include <glog/logging.h>

#include <QtGlobal>

#include "octopod.h"
#include "tcpsocket.h"

namespace octopod {

// Creates a TCP transport socket, initially disconnected.
TCPSocket::TCPSocket() : sock_(BAD_SOCKET), state_(kDisconnected), ipversion_(NetworkAddress::kIPv4) { }

// Cleans up the socket structure.
// Does not close the connection explicitly.
TCPSocket::~TCPSocket() {
  if (sock_ != BAD_SOCKET) {
    disconnect();
  }
}

// Returns the current state by asking the underlying Qt socket.
// Currently doesn't distinguish between several states that should
// eventually be treated differrently.
Socket::SocketState TCPSocket::state() const {
  return state_;
}

// Connects to the given address.  Does not support NAT hole punching
// at the moment, so the NAT portion of the address is ignored.
// Also currently waits forever for a connection to be established.
Socket::SocketError TCPSocket::connect(const NetworkAddress &addr) {
  // Sanity check the given address
  if (addr.hostaddr().empty()) {
    DLOG(WARNING) << Q_FUNC_INFO << "Given empty host address to connect to.";
    return kFail;
  }

  if (!addr.nataddr().empty()) {
    DLOG(WARNING) << Q_FUNC_INFO << "TCPSocket doesn't support NAT traversal.";
    return kFail;
  }

  struct sockaddr *sa = NULL;
  int sa_len = -1;
  struct sockaddr_in sa4;
  struct sockaddr_in6 sa6;

  // Determine whether it's IPv4 or IPv6 and parse.
  if (addr.version() == NetworkAddress::kIPv4) {
    memset(&sa4, 0, sizeof(sa4));

    int err = inet_pton(AF_INET, addr.hostaddr().c_str(), &sa4.sin_addr);
    if (err == 0) {
      DLOG(WARNING) << Q_FUNC_INFO << "Address " << addr.hostaddr().c_str() << " wasn't a valid IPv4 address.";
      return kFail;
    } else if (err == -1) {
      DLOG(WARNING) << Q_FUNC_INFO << "Address " << addr.hostaddr().c_str() << " couldn't be parsed by inet_pton(): " << strerror(errno);
      return kFail;
    }

    sa4.sin_family = AF_INET;
    sa4.sin_port = htons(addr.port());

    sa = (struct sockaddr *) &sa4;
    sa_len = sizeof(sa4);
  } else if (addr.version() == NetworkAddress::kIPv6) {
    memset(&sa6, 0, sizeof(sa6));
    
    int err = inet_pton(AF_INET6, addr.hostaddr().c_str(), &sa6.sin6_addr);
    if (err == 0) {
      DLOG(WARNING) << Q_FUNC_INFO << "Address " << addr.hostaddr().c_str() << " wasn't a valid IPv6 address.";
      return kFail;
    } else if (err == -1) {
      DLOG(WARNING) << Q_FUNC_INFO << "Address " << addr.hostaddr().c_str() << " couldn't be parsed by inet_pton(): " << strerror(errno);
      return kFail;
    }

    sa6.sin6_family = AF_INET6;
    sa6.sin6_port = htons(addr.port());

    sa = (struct sockaddr *) &sa6;
    sa_len = sizeof(sa6);
  } else {
    DLOG(WARNING) << Q_FUNC_INFO << "Address is of unknown version.";
    return kFail;
  }
  
  if (sock_ != BAD_SOCKET && state_ != kDisconnected) {
    DLOG(WARNING) << Q_FUNC_INFO << "Asked to connect using an already in use socket.";
    return kFail;
  }

  // If the socket exists but is for the wrong IP version, close it.
  if (sock_ != BAD_SOCKET && ipversion_ != addr.version()) {
    closesocket(sock_);
    sock_ = BAD_SOCKET;
  }

  // Create a new socket if needed.
  if (sock_ == BAD_SOCKET) {
    int domain = addr.version();
    sock_ = socket(domain, SOCK_STREAM, 0); 
    if (sock_ == BAD_SOCKET) {
      DLOG(WARNING) << Q_FUNC_INFO << "socket() failed with error " << strerror(errno);
      return kFail;
    }
  }

  // Connect to the remote host.
  if (::connect(sock_, sa, sa_len) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "connect() failed with error " << strerror(errno);
    return kFail;
  }

  state_ = kConnected;

  return kOk;
}

// Disconnects the current connection.  Blocks until any buffered data
// is sent to the underlying socket before returning.
Socket::SocketError TCPSocket::disconnect() {
  if (sock_ == BAD_SOCKET || state_ != kConnected) {
    return kFail;
  }

  if (closesocket(sock_) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "couldn't close socket, failed with error " << strerror(errno);
    return kFail;
  }

  state_ = kDisconnected;

  return kOk;
}

// Sends up to size bytes from data.
int TCPSocket::send(const char *data, int size) {
  CHECK_NOTNULL(data);
  CHECK(size >= 0);

  if (sock_ == BAD_SOCKET || state_ != kConnected) {
    DLOG(WARNING) << Q_FUNC_INFO << "shouldn't have been called when disconnected.";
    return -1;
  }

  int ret;
  if ((ret = ::send(sock_, data, size, 0)) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "send() failed with error " << strerror(errno);
    return -1;
  }

  return ret; 
}

// Reads up to size bytes
int TCPSocket::recv(int size, char *data) {
  CHECK_NOTNULL(data);
  CHECK(size >= 0);

  if (sock_ == BAD_SOCKET || state_ != kConnected) {
    DLOG(WARNING) << Q_FUNC_INFO << "shouldn't have been called when disconnected.";
    return -1;
  }
  
  int ret;
  if ((ret = ::recv(sock_, data, size, 0)) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "send() failed with error " << strerror(errno);
    return -1;
  }

  return ret;
}

TCPListenSocket::TCPListenSocket() {
  listen_sock_ = BAD_SOCKET;
  state_ = kDisconnected;
}

TCPListenSocket::~TCPListenSocket() {
  closesocket(listen_sock_);
}

// Starts listening on address and port given.
// If addr.hostaddr() is empty, binds to INADDR_ANY
// If addr.port() is 0, a port is chosen by the underlying implementation
Socket::SocketError TCPListenSocket::listen(uint16_t port) {
  if (state_ != kDisconnected) {
    return Socket::kFail;
  }

  // Create socket.
  listen_sock_ = socket(AF_INET6, SOCK_STREAM, 0); 
  if (listen_sock_ == BAD_SOCKET) {
    DLOG(WARNING) << Q_FUNC_INFO << "socket() failed with error " << strerror(errno);
    return Socket::kFail;
  }

  // Allow reuse of recently freed port.
  const int yes = 1;

  if (setsockopt(listen_sock_, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes)) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "Unable to set SO_REUSEADDR on socket: " << strerror(errno);
    closesocket(listen_sock_);
    return Socket::kFail;
  }

  // Allow use of the socket for both IPv6 and IPv4.
  const int no = 0;

  if (setsockopt(listen_sock_, IPPROTO_IPV6, IPV6_V6ONLY, &no, sizeof(no)) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "Unable to unset IPV6_V6ONLY on socket: " << strerror(errno);
    closesocket(listen_sock_);
    return Socket::kFail;
  }

  // Bind to the given port.
  struct sockaddr_in6 sa6;
  memset(&sa6, 0, sizeof(sa6));

  sa6.sin6_family = AF_INET6;
  sa6.sin6_port = htons(port);
  sa6.sin6_addr = in6addr_any;
  sa6.sin6_len = sizeof(sa6);

  if (bind(listen_sock_, (struct sockaddr *) &sa6, sizeof(sa6)) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "Unable to bind socket: " << strerror(errno);
    closesocket(listen_sock_);
    return Socket::kFail;
  }

  // Start listening.
  if (::listen(listen_sock_, kIncomingQueueLength) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "Unable to listen(): " << strerror(errno);
    closesocket(listen_sock_);
    return Socket::kFail;
  }

  state_ = kListening;

  return Socket::kOk;
}

// Blocks for a new incoming TCP connection. 
// Returns NULL only upon error.
Socket *TCPListenSocket::accept() {
  if (state_ != kListening) {
    DLOG(WARNING) << Q_FUNC_INFO << "Listening socket is not in listening state.";
    return NULL;
  }

  // Block forever for an incoming TCP connection.
  sockaddr_storage sa;
  int sa_len = sizeof(sa);

  int sock;
  if ((sock = ::accept(listen_sock_, (struct sockaddr *) &sa, (socklen_t *) &sa_len)) == SOCKET_ERROR) {
    DLOG(WARNING) << Q_FUNC_INFO << "Unable to accept(): " << strerror(errno);
    return NULL;
  }

  // Print info about the accepted connection.
  char str[INET6_ADDRSTRLEN];
  int port = 0;
  memset(str, 0, sizeof(str));

  if (sa.ss_family == AF_INET) {
    inet_ntop(AF_INET, &(((struct sockaddr_in *) &sa)->sin_addr), str, sizeof(str));
    port = ntohs(((struct sockaddr_in *) &sa)->sin_port);
  } else if (sa.ss_family == AF_INET6) {
    inet_ntop(AF_INET6, &(((struct sockaddr_in6 *) &sa)->sin6_addr), str, sizeof(str));
    port = ntohs(((struct sockaddr_in6 *) &sa)->sin6_port);
  } else {
    DLOG(WARNING) << Q_FUNC_INFO << "Got an incoming connection of unknown family " << sa.ss_family;
    closesocket(sock);
    return NULL;
  }
  
  LOG(INFO) << Q_FUNC_INFO << "Accepted a connection from " << str << ":" << port;

  // Initialize a new socket object with the descriptor.
  TCPSocket *s = new TCPSocket();
  s->sock_ = sock;
  s->state_ = TCPSocket::kConnected;
  s->ipversion_ = (sa.ss_family == AF_INET) ? NetworkAddress::kIPv4 : NetworkAddress::kIPv6;

  return s;
}

}  // namespace octopod
